% 定性展示氢原子的s轨道
% R的函数参考另一个脚本
% Gitee Repo

clc
clear

function R = Rnl(r, a, n, l)
  kappa = 1/(a*n);
  rho = kappa * r;

  R = 1;
  c = zeros(1,n);

  for j = 1:n-l-1;
    C = 2*(j+l-n)/(j*(j+2*l+1));
    if j == 1
      c(j) = C*1;
    else
      c(j) = C*c(j-1);
    end
    R += c(j)*(rho.^j);
  end
  R = (1./r).*rho.^(l+1).*(e.^(-rho)).*R;
end


a = 0.2;
n = 2;

% 根据AI的建议优化为矩阵形式
[x y] = meshgrid(-5:0.05:5);
r = sqrt(x.^2+y.^2);
Psi = Rnl(r, a, n, 0).^2*10;
invalid_values = isnan(Psi) | isinf(Psi) | (Psi > 1);
Psi(invalid_values) = 1;

% for的过慢
##sx = size(x,1);
##for i = 1:sx
##  for j = 1:sx
##    r = sqrt(x(i,j)^2+y(i,j)^2);
##    Psi(i,j) = Rnl(r, a, n, 0);
##    Psi(i,j) = Psi(i,j).^2*10;
##    if isnan(Psi(i,j)) || isinf(Psi(i,j)) || Psi(i,j) > 1
##      Psi(i,j)  = 1;
##    end
##  end
##end


figure();
hold on
axis equal
h = surf(x,y,Psi);
set(h,'edgecolor','none')
